

open Ast
open Format

module StrMap = Map.Make(String)

(*
  TODO : typage

*)


(* les differentes erreurs *)

type var_typing_error = 
  | VTE_unbound of typ
  | VTE_void 
  | VTE_array_funtype 
      
type error = 
  | VarAlreadyExist of string
  | Unbound_var of string
  | Unbound_array of string
  | Unbound_value of string
  | Enum_value_already_exists of string
  | Var_typing_error of var_typing_error * string
  | TypeAlreadyExist of typ
  | Bad_assign 
  | Invalid_compop of compop * typ * typ
  | Invalid_boolop of logicop * typ * typ
  | Invalid_arithop of arithop * typ * typ
  | Not_compatible of typ * typ
  | Not_assign_in_atomic
  | Await_expr
  | Await_assign_or_atomic
  | Wrong_process_order

let rec print_err_type fmt = function
| Tint -> eprintf "int"
| Treal -> eprintf "real"
| Tproc -> eprintf "proc"
| Tbool -> eprintf "bool"
| Tident id -> eprintf "%s" id.node
| Tproccnt -> eprintf "proccnt"
| Tany -> eprintf "any"
| Tlinenum _ -> eprintf "line"


let print_err_compop fmt = function 
| Beq -> eprintf "binary =="
| Bne -> eprintf "binary !="
| Blt -> eprintf "binary <"
| Ble -> eprintf "binary <="
| Bgt -> eprintf "binary >"
| Bge -> eprintf "binary >="


let print_err_arithop fmt = function 
| Bplus -> eprintf "binary +"
| Bminus -> eprintf "binary -"

    
let print_err_boolop fmt = function 
| Band -> eprintf "binary &&"
| Bor -> eprintf "binary ||"

    
let print_var_typing_error fmt = function
| VTE_unbound ct -> eprintf "Unbound '%a'" 
    print_err_type ct
| VTE_void -> eprintf "Forbidden type 'void'"
| VTE_array_funtype -> eprintf "Forbidden type array"

    
let print_error = function
| Await_expr -> 
    eprintf "Error : this expression must be a quantifier"
| VarAlreadyExist s -> 
    eprintf "Error : redefinition of the variable '%s'\n@." s
| TypeAlreadyExist ct -> 
    eprintf 
      "Error : redefinition of '%a'\n@." 
      print_err_type ct
| Var_typing_error (vte, s) ->
    eprintf 
      "Error : %a in the definition of the variable '%s'\n@."
      print_var_typing_error vte s
| Unbound_var s -> 
    eprintf "Error : Unbound variable '%s'\n@." s
| Unbound_value s -> 
    eprintf "Error : Unbound value '%s'\n@." s
| Enum_value_already_exists s -> 
    eprintf "Error : this enum value alreay exists '%s'\n@." s
| Unbound_array s -> 
    eprintf "Error : Unbound array '%s'\n@." s
| Invalid_compop (op, t1, t2) -> 
    eprintf "Error : invalid operand to %a ('%a' and '%a')\n@." 
      print_err_compop op print_err_type t1 print_err_type t2 
| Invalid_boolop (op, t1, t2) -> 
    eprintf "Error : invalid operand to %a ('%a' and '%a')\n@." 
      print_err_boolop op print_err_type t1 print_err_type t2 
| Invalid_arithop (op, t1, t2) -> 
    eprintf "Error : invalid operand to %a ('%a' and '%a')\n@." 
      print_err_arithop op print_err_type t1 print_err_type t2 
| Bad_assign -> eprintf "Error : The left value cannot be assigned\n@." 
| Not_assign_in_atomic ->
    eprintf
      "Error : only assign statement is allowed inside an atomic block\n@." 
| Not_compatible (t1, t2)-> 
    eprintf "Error : type '%a' is not compatible with '%a'\n@." 
      print_err_type t1 print_err_type t2
| Await_assign_or_atomic ->
    eprintf
      "Error : only assign and atomic statement are allowed inside await\n@." 
| Wrong_process_order ->
    eprintf
      "Error : process must be defined in the right order (starting by 1)\n@." 


(* Exception generique utilisee au typage *)
exception Error of error * Ast.loc

(* exception particulière concernant l'absence de main *)
exception No_main_error 

(* exception utilisee et rattrape localement *)
exception Incompatible

type kind = Sconst | Sglobal | Slocal | Sarray  

let types = Hashtbl.create 32 (* (id, typ) *)
let enums = Hashtbl.create 17 (* idenum, type *)
let vars = ref StrMap.empty (* (id, (kind, typ)) *)

let is_global s = 
  try
    begin
      let k, t = StrMap.find s !vars in
      match k with
      | Sglobal | Sconst | Sarray -> true | Slocal -> false
    end
  with | Not_found -> false
  
let clear_envs () =
  Hashtbl.clear types;
  vars := StrMap.empty
    
let mk_node e t = {data = t; node = e}
  
let loc_dummy e =
  mk_node e (Lexing.dummy_pos, Lexing.dummy_pos)

let rec type_stmt env stmt =
  match stmt.node with
  | Sawait (qt, so) ->
      let tqt = type_qterm env qt in
      (*      begin match tqt.node with
              | Qquantifier _ -> () | _ -> raise (Error (Await_expr, qt.data)) end; *)
      (* In the end, we accept await <logical_term> *)
      let tso = begin match so with
      | None -> None
      | Some s -> begin match s.node with
        | Sassign _ | Satomic _ -> Some (type_stmt env s)
        | _ -> raise (Error (Await_assign_or_atomic, s.data))
      end end in
      mk_node (Sawait (tqt, tso)) (Tlinenum stmt.data)
  | Sassign (lv, term) ->
      let tlv = type_leftval env lv in
      let tterm = type_term env term in
      if not (tlv.data = tterm.data) then
      raise (Error (Not_compatible (tlv.data, tterm.data), lv.data));
      mk_node (Sassign (tlv, tterm)) (Tlinenum stmt.data)
  | Sifelse (qt, sl1, sl2) ->
      let tqt = type_qterm env qt in
      let tsl1 = type_stmts env sl1 in
      let tsl2 = type_stmts env sl2 in
      if not (tqt.data = Tbool) then
      raise (Error (Not_compatible (tqt.data, Tbool), qt.data));
      mk_node (Sifelse (tqt, tsl1, tsl2)) (Tlinenum stmt.data) 
  | Scritical -> mk_node Scritical (Tlinenum stmt.data)
  | Swhile (qt, sl) ->
      let tqt = type_qterm env qt in
      let tsl = type_stmts env sl in
      mk_node (Swhile (tqt, tsl)) (Tlinenum stmt.data)
  | Sassert qt -> mk_node (Sassert (type_qterm env qt)) (Tlinenum stmt.data)
  | Satomic stmts ->
      let tstmts = 
        List.map (fun stmt ->
          begin match stmt.node with
          | Sassign _ -> type_stmt env stmt
          | _ -> raise (Error (Not_assign_in_atomic, stmt.data))
          end) stmts
      in
      mk_node (Satomic tstmts) (Tlinenum stmt.data)

and type_leftval env lv =
  match lv.node with
  | LVident id ->
      let _, t = begin try StrMap.find id.node env
        with | Not_found -> raise (Error (Unbound_var id.node, id.data))
      end in mk_node (LVident id) t
  | LVaccess (id, ids) ->
      let ka, ta =
        begin try StrMap.find id.node env with
        | Not_found -> raise (Error (Unbound_array id.node, id.data))
        end
      in
      List.iter (fun index ->
        let ki, ti =
          begin try StrMap.find index.node env with
          | Not_found -> raise (Error (Unbound_var index.node, index.data))
          end
        in
        if ((ki <> Slocal && ki <> Sglobal) || ti <> Tproc) then
        raise (Error (Not_compatible (ti, Tproc), index.data))
      ) ids;
      if ka <> Sarray then raise (Error (Unbound_array id.node, id.data))
      else mk_node (LVaccess (id, ids)) ta

and type_term env t =
  match t.node with
  | TMatom at ->
      let tat = type_atomic_term env at in
      mk_node (TMatom tat) tat.data
  | TMarith (op, at, c) ->
      let tat = type_atomic_term env at in
      let tc = type_constnum c in
      let t = begin match tat.data, tc.data with
      | (Tint, Tint) -> Tint
      | Treal, Treal -> Treal
      | Treal, Tint | Tint, Treal ->
          raise (Error ((Not_compatible (tat.data, tc.data), at.data)))
      | _ ->
          raise (Error ((Invalid_arithop (op, tat.data, tc.data),
                         at.data)))
      end in
      mk_node (TMarith (op, tat, tc)) t

and type_constnum c =
  match c.node with
  | Creal f -> mk_node (Creal f) Treal 
  | Cint i -> mk_node (Cint i) Tint
      
and type_atomic_term env at =
  match at.node with
  | ATleftval lv ->
      let tlv = type_leftval env lv in 
      mk_node (ATleftval tlv) tlv.data 
  | ATconstnum const ->
      let tconst = type_constnum const in
      mk_node (ATconstnum tconst) tconst.data
  | ATbool b -> 
      mk_node (ATbool b) Tbool
  | ATenum id -> 
      try  
      mk_node (ATenum id) (Tident (Hashtbl.find enums id.node))
      with Not_found -> raise (Error (Unbound_value id.node, id.data))
        
and type_stmts env sl =
  List.map (type_stmt env) sl

and type_quantifier env q id e = 
  let env' = StrMap.add id.node (Slocal, Tproc) env in
  mk_node (Qquantifier (q, id, type_logic_term env' e)) Tbool

and type_logic_term env lt =
  match lt.node with
  | LTterm t ->
      mk_node (LTterm (type_term env t)) Tbool
  | LTlit lit -> mk_node (LTlit (type_literal env lit)) Tbool
  | LTop (op, lt1, lt2) ->
      let tlt1, tlt2 = type_logic_term env lt1, type_logic_term env lt2 in
      mk_node (LTop (op, tlt1, tlt2)) Tbool

and type_literal env lit =
  let op, term1, term2 = lit.node in
  let tt1, tt2 = type_term env term1, type_term env term2 in
  mk_node (op, tt1, tt2) Tbool

and type_qterm env qe =
  match qe.node with
  | Qquantifier (q, id, e) ->
      type_quantifier env q id e
  | Qterm lt ->
      let tlt = type_logic_term env lt in
      mk_node (Qterm tlt) Tbool

and type_type_decl (idt, idl) =
  let ids =
    List.fold_left (fun ids idenum ->
      if Hashtbl.mem enums idenum.node then
      raise (Error ((Enum_value_already_exists idenum.node), idenum.data))
      else Hashtbl.add enums idenum.node idt;
      idenum.node :: ids) [] idl in
  Hashtbl.add types idt.node (List.rev ids);
  (idt, idl)

and type_valid var t = match t with
| Tident id -> if Hashtbl.mem types id.node then () else
  raise (Error (Var_typing_error((VTE_unbound t), var.node), var.data))
| _ -> ()

and type_var_decl = function
| VDconst (id, t) as d -> type_valid id t;
    if StrMap.mem id.node !vars
    then raise (Error (VarAlreadyExist id.node, id.data)) 
    else vars := StrMap.add id.node (Sconst, t) !vars; d
| VDglobal (id, t) as d -> type_valid id t;
    if StrMap.mem id.node !vars
    then raise (Error (VarAlreadyExist id.node, id.data)) 
    else vars := StrMap.add id.node (Sglobal, t) !vars;
    d
| VDarray (id, it, t) as d -> type_valid id t;
    if StrMap.mem id.node !vars
    then raise (Error (VarAlreadyExist id.node, id.data)) 
    else vars := StrMap.add id.node (Sarray, t) !vars; d

and type_decls (dl : (Ast.loc Ast.decl) list) : (Ast.typ Ast.decl) list =
  List.fold_left(
    fun (dl) d -> (
      match d with
      | Dtype td -> Dtype (type_type_decl td)
      | Dvar vd -> Dvar (type_var_decl vd)
    ) :: dl) [] dl

    
let type_procs pl =
  let ids, stmtsl = List.split pl in 
  ignore(
    List.fold_left (fun acc ident ->
      let prec = acc in 
      let id = int_of_string
        (String.sub ident.node 1 (String.length ident.node - 1)) in
      if id <> prec + 1 then (raise (Error (Wrong_process_order, ident.data)));
      vars := StrMap.add ident.node (Sglobal, Tproc) !vars; id) 0 ids);
  List.combine ids (List.map (type_stmts !vars) stmtsl)
    
    
let type_prog (p : Ast.loc Ast.prog) : Ast.typ Ast.prog =
  let decls, procs = p in
  let tdecls = type_decls decls in
  let tprocs =
    match procs with
    | [] -> []
    | [cp, stmts] -> 
        vars := StrMap.add cp.node (Slocal, Tproc) !vars;
        [cp, type_stmts !vars stmts]
    | l -> type_procs l
  in
  (tdecls, tprocs)


















